
import insightface
import numpy as np
import cv2
from uniform_cls import ONNXClassifier

class PersonDet():

    def __init__(self):
        self.bdet = insightface.model_zoo.get_model('scrfd_person_2.5g.onnx', download=True)
        self.bdet.prepare(0, det_thresh=0.4, nms_thresh=0.6, input_size=(640, 640))

    def body_detect(self, image: np.ndarray):
        bboxes, kpss = self.bdet.detect(image)
        return self.detect_person(image, bboxes, kpss)

    @staticmethod
    def detect_person(img: np.ndarray, bboxes, kpss):
        bboxes = np.round(bboxes[:, :4]).astype(np.int)
        kpss = np.round(kpss).astype(np.int)
        kpss[:, :, 0] = np.clip(kpss[:, :, 0], 0, img.shape[1])
        kpss[:, :, 1] = np.clip(kpss[:, :, 1], 0, img.shape[0])
        vbboxes = bboxes.copy()
        vbboxes[:, 0] = kpss[:, 0, 0]
        vbboxes[:, 1] = kpss[:, 0, 1]
        vbboxes[:, 2] = kpss[:, 4, 0]
        vbboxes[:, 3] = kpss[:, 4, 1]
        return bboxes, vbboxes

    @staticmethod
    def draw(img: np.ndarray, bboxes, vbboxes, res, colors=None):
        if colors is None:
            colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0)]  # 默认颜色
        labels = ["stranger", "employee", "manager"]

        for i in range(bboxes.shape[0]):
            bbox = bboxes[i]
            vbbox = vbboxes[i]
            x1, y1, x2, y2 = bbox
            vx1, vy1, vx2, vy2 = vbbox
            cv2.rectangle(img, (x1, y1), (x2, y2), colors[res[i]], 1)

            # 添加文本
            label = labels[res[i]]
            font_scale = 0.5  # 默认字体大小
            text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1)[0]
            text_width = text_size[0]
            box_width = x2 - x1

            # 如果文本宽度大于框宽度，则调整字体大小
            if text_width > box_width:
                font_scale *= box_width / text_width

            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, colors[res[i]], 1)

            # 绘制虚拟框
            alpha = 0.8
            color = colors[res[i]]
            for c in range(3):
                img[vy1:vy2, vx1:vx2, c] = img[vy1:vy2, vx1:vx2, c] * alpha + color[c] * (1.0 - alpha)

            # 绘制关键点
            cv2.circle(img, (vx1, vy1), 1, color, 2)
            cv2.circle(img, (vx1, vy2), 1, color, 2)
            cv2.circle(img, (vx2, vy1), 1, color, 2)
            cv2.circle(img, (vx2, vy2), 1, color, 2)

        return img



    def process_video(self, video_path, is_show=False):
        cap = cv2.VideoCapture(video_path)

        # 获取视频的基本属性
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)

        # 设置输出视频的格式和名称（保存为 MP4）
        output_video_path = video_path.rsplit(".", 1)[0] + "_uniform.mp4"
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 或者 'avc1'
        out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # 应用人体检测和分类逻辑
            bboxes, vbboxes = self.body_detect(frame)
            res = []
            for i in range(bboxes.shape[0]):
                bbox = bboxes[i]
                x1, y1, x2, y2 = bbox
                cropped = frame[y1:y2, x1:x2]
                cls_res = cls.predict(cropped)
                res.append(cls_res)  # 假设 cls.predict 返回一个单元素数组

            self.draw(frame, bboxes, vbboxes, res)

            # 将处理后的帧写入输出文件
            out.write(frame)

            # 根据 is_show 决定是否显示帧
            if is_show:
                cv2.imshow('Frame', frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        # 释放资源
        cap.release()
        out.release()
        if is_show:
            cv2.destroyAllWindows()

def image_test():
    cls = ONNXClassifier("model.onnx")
    det = PersonDet()
    img = cv2.imread("/Users/tunm/Downloads/WechatIMG1679.jpg")

    bboxes, vbboxes = det.body_detect(img)
    bboxes = np.round(bboxes[:, :4]).astype(np.int)
    for i in range(bboxes.shape[0]):
        bbox = bboxes[i]
        x1, y1, x2, y2 = bbox
        cropped = img[y1:y2, x1:x2]
        res = cls.predict(cropped)
        print(res)
        cv2.imshow("cropped", cropped)
        cv2.waitKey(0)


if __name__ == "__main__":
    cls = ONNXClassifier("model.onnx")
    det = PersonDet()
    det.process_video("/Users/tunm/Downloads/进场+工位视频素材/圈出【员工身体】进入【工位】+圈出【员工脸部】+【员工拿出手机】打电话.MOV")  # 替换为你的视频文件路径
